#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys, os
import moveit_commander
from moveit_msgs.msg import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
from geometry_msgs.msg import PoseStamped, Pose, PointStamped
from vision_msgs.msg import Detection2D
from std_msgs.msg import Bool
import tf2_ros
import tf2_geometry_msgs
import math
import random

class MoveItGrasp:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('moveit_grasp_demo')
        r = rospy.Rate(10)
        rospy.on_shutdown(self.shutdown)
                
        # 初始化MoveIt命令器
        moveit_commander.roscpp_initialize(sys.argv) 
        self.arm = moveit_commander.MoveGroupCommander('manipulator')        
        self.end_effector_link = self.arm.get_end_effector_link()
        self.reference_frame = 'base_link'
        self.arm.set_pose_reference_frame(self.reference_frame)

        # 初始化TF2
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        # MoveIt规划参数
        self.arm.set_planning_time(10.0)
        self.arm.set_num_planning_attempts(50)
        self.arm.set_max_velocity_scaling_factor(0.3)
        self.arm.set_max_acceleration_scaling_factor(0.3)
        self.arm.allow_replanning(True)
        self.arm.set_goal_position_tolerance(0.01)
        self.arm.set_goal_orientation_tolerance(0.05)
        self.arm.set_planner_id("RRTConnect")  # 指定规划器
        rospy.loginfo("Planner ID set to: %s", self.arm.get_planner_id())

        # 抓取控制变量
        self.cnt = 0
        self.stable_cnt = 40  # 稳定计数，确保目标稳定
        self.grasp_status = False
        self.object_list = []  # 存储所有检测到的物体
        self.grasped_positions = set()  # 存储已抓取的物体位置（3D）
        self.distance_threshold = rospy.get_param('~distance_threshold', 0.001)  # 距离阈值（米），默认 0.02
        self.max_grasp_per_cycle = rospy.get_param('~max_grasp_per_cycle', 3)    # 每次循环最多抓取的物体数量，默认3

        # 夹爪偏移量（从末端执行器到抓取点，单位：米）
        self.gripper_offset_x = rospy.get_param('~gripper_offset_x', 0.085)
        self.gripper_offset_y = rospy.get_param('~gripper_offset_y', -0.02)
        self.gripper_offset_z = rospy.get_param('~gripper_offset_z', -0.0005)

        # 订阅对象坐标
        self.object_sub = rospy.Subscriber("/object_point_info", Detection2D, self.ObjectCallback)

        # 初始化发布器
        self.grasped_pub = rospy.Publisher('/grasped_positions', PointStamped, queue_size=1)
        self.reset_pub = rospy.Publisher('/reset_grasped', Bool, queue_size=1)

        # 移动到初始姿态
        self.init_pose()
        rospy.sleep(1)

        """
            将self.gripper_open()与self.gripper_close()放置在理想位置以实现抓取逻辑
            self.gripper_open()中设置了两种，可按需保留，一种方便调参
            力度设置按需保留

            open与close区别在于数值，即开合角度距离，代码基本一致，无本质区别
        """
        # os.system("rosservice call /io_service/set_gripper_position \"{val: 90}\"")
        # self.gripper_open(10)
        # rospy.sleep(1)
            

        while not rospy.is_shutdown():
            if self.grasp_status and len(self.object_list) > 0:
                # 获取最多max_grasp_per_cycle个未抓取的目标（随机选择）
                target_objs = self.get_next_ungrasped_objects()
                if not target_objs:
                    rospy.loginfo("All targets grasped or no valid targets (grasped_positions=%s)", str(self.grasped_positions))
                    self.check_and_reset_grasp()
                    self.grasp_status = False
                    self.object_list = []
                    self.cnt = 0
                    rospy.sleep(2.5)  # 延迟以等待视觉系统更新
                    continue

                # 依次抓取选中的目标
                for target_obj in target_objs:
                    rospy.loginfo("Starting grasp for tomato: x=%.3f, y=%.3f, z=%.3f", 
                                  target_obj.position.x, target_obj.position.y, target_obj.position.z)
                    x, y, z = self.camera_to_gripper_pose(target_obj.position.x, target_obj.position.y, target_obj.position.z)
                    if any(v is None for v in (x, y, z)):
                        rospy.logerr("Invalid grasp coordinates: x=%s, y=%s, z=%s", x, y, z)
                        continue
                    rospy.loginfo("Moving to grasp position: %f, %f, %f", x, y, z)
                    if self.move_pose(x, y, z):
                        # 记录已抓取的位置
                        self.grasped_positions.add((round(x, 3), round(y, 3), round(z, 3)))
                        rospy.loginfo("Grasped tomato at: x=%.3f, y=%.3f, z=%.3f", x, y, z)
                        # 发布已抓取位置
                        self.publish_grasped_position(target_obj)
                    else:
                        rospy.logwarn("Grasp attempt failed, skipping position recording")
                    self.check_and_reset_grasp()
                    rospy.sleep(1)
                    self.init_pose()
                    rospy.sleep(2)
                self.grasp_status = False
                self.object_list = []
                self.cnt = 0
                rospy.sleep(2.5)  # 额外延迟以等待视觉系统更新
            r.sleep()

    def ObjectCallback(self, msg):
        """处理物体检测消息"""
        obj_info = msg
        self.cnt += 1
        if self.cnt > self.stable_cnt and not self.grasp_status:
            self.grasp_status = True
            # 收集所有检测到的物体
            self.object_list = [result.pose.pose for result in obj_info.results]  # 包含 position 和 orientation
            if len(self.object_list) == 1:
                rospy.loginfo("Detected single tomato: x=%.3f, y=%.3f, z=%.3f", 
                              self.object_list[0].position.x, self.object_list[0].position.y, self.object_list[0].position.z)
            else:
                rospy.loginfo("Detected %d tomatoes", len(self.object_list))

    def get_next_ungrasped_objects(self):
        """获取最多max_grasp_per_cycle个未抓取的物体（随机选择）"""
        ungrasped_objs = []
        for obj in self.object_list:
            x, y, z = self.camera_to_gripper_pose(obj.position.x, obj.position.y, obj.position.z)
            if any(v is None for v in (x, y, z)):
                rospy.logwarn("Skipping object due to invalid coordinates")
                continue
            # 检查是否接近已抓取的位置
            is_grasped = False
            for grasped_x, grasped_y, grasped_z in self.grasped_positions:
                distance = math.sqrt((x - grasped_x)**2 + (y - grasped_y)**2 + (z - grasped_z)**2)
                if distance < self.distance_threshold:
                    rospy.loginfo("Tomato at x=%.3f, y=%.3f, z=%.3f filtered as grasped (distance=%.3f, grasped=%s)", 
                                  x, y, z, distance, str((grasped_x, grasped_y, grasped_z)))
                    is_grasped = True
                    break
            if not is_grasped:
                ungrasped_objs.append(obj)

        if not ungrasped_objs:
            rospy.loginfo("No ungrasped tomatoes found")
            return []

        # 随机选择最多max_grasp_per_cycle个目标
        selected_objs = random.sample(ungrasped_objs, min(self.max_grasp_per_cycle, len(ungrasped_objs)))
        rospy.loginfo("Selected %d ungrasped tomatoes for grasp", len(selected_objs))
        for obj in selected_objs:
            x, y, z = self.camera_to_gripper_pose(obj.position.x, obj.position.y, obj.position.z)
            rospy.loginfo("Selected tomato: x=%.3f, y=%.3f, z=%.3f", x, y, z)
        return selected_objs

    def check_and_reset_grasp(self):
        """检查是否还有未抓取目标，若无则清空记录"""
        if not self.get_next_ungrasped_objects():
            rospy.loginfo("No more ungrasped targets, resetting grasped positions")
            self.grasped_positions.clear()
            # 发布重置信号
            reset_msg = Bool()
            reset_msg.data = True
            self.reset_pub.publish(reset_msg)
            rospy.loginfo("Published reset signal")

    def publish_grasped_position(self, obj):
        """发布已抓取位置（相机框架的像素坐标）"""
        try:
            grasped_point = PointStamped()
            grasped_point.header.frame_id = "camera_link"
            grasped_point.header.stamp = rospy.Time.now()
            grasped_point.point.x = obj.orientation.x  # 使用 object_xyz.py 传递的像素坐标
            grasped_point.point.y = obj.orientation.y
            grasped_point.point.z = 0  # 仅使用 x, y
            self.grasped_pub.publish(grasped_point)
            rospy.loginfo("Published grasped position: x=%.1f, y=%.1f", 
                          grasped_point.point.x, grasped_point.point.y)
        except Exception as e:
            rospy.logerr("Failed to publish grasped position: %s", str(e))

    def camera_to_gripper_pose(self, x, y, z):
        """
        将对象坐标（毫米，相机框架）转换为夹爪姿态（米，base_link框架）。
        """
        point_camera = PointStamped()
        point_camera.header.frame_id = "camera_link"
        point_camera.header.stamp = rospy.Time(0)
        point_camera.point.x = x / 1000.0
        point_camera.point.y = y / 1000.0
        point_camera.point.z = z / 1000.0

        try:
            point_base = self.tf_buffer.transform(point_camera, "world", 
                                                timeout=rospy.Duration(2.0))
            grasp_x = point_base.point.x
            grasp_y = point_base.point.y
            grasp_z = point_base.point.z
            rospy.loginfo("Transformed tomato point: x=%.3f, y=%.3f, z=%.3f", 
                          grasp_x, grasp_y, grasp_z)
            return grasp_x, grasp_y, grasp_z
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, 
                tf2_ros.ExtrapolationException) as e:
            rospy.logerr("TF transform failed: %s", str(e))
            return None, None, None

    def move_pose(self, x, y, z):
        """
        移动机械臂到指定位置
        返回：bool，是否成功执行
        """
        target_pose = PoseStamped()
        target_pose.header.frame_id = self.reference_frame
        target_pose.header.stamp = rospy.Time.now()     
        current_pose = self.arm.get_current_pose(self.end_effector_link)
        target_pose.pose.position.x = x + self.gripper_offset_x
        target_pose.pose.position.y = y + self.gripper_offset_y 
        target_pose.pose.position.z = z + self.gripper_offset_z
        target_pose.pose.orientation = current_pose.pose.orientation
        self.arm.set_start_state_to_current_state()
        self.arm.set_pose_target(target_pose, self.end_effector_link)
        
        traj = self.arm.plan()
        if traj and (hasattr(traj, 'joint_trajectory') and traj.joint_trajectory.points):
            success = self.arm.execute(traj, wait=True)
            if success:
                rospy.loginfo("Successfully executed grasp motion")
                return True
            else:
                rospy.logwarn("Execution failed")
                return False
        else:
            rospy.logwarn("Planning failed: No valid trajectory found")
            return False

    def init_pose(self):
        """移动到初始观察状态"""
        self.arm.set_named_target("look_L")
        success = self.arm.go(wait=True)
        if success:
            rospy.loginfo("Successfully moved to look state")
        else:
            rospy.logwarn("Failed to move to look state")
    
    def gripper_open(self, value):
        rospy.loginfo("Start open gripper")
        os.system("rosservice call /io_service/set_gripper_position \"{val: %s}\""%value)  # lebai原夹爪
        # os.system("rosservice call /io_service/set_gripper_force \"{val: 10}\"")    # lebai夹爪抓力，可选
        # os.system("rostopic pub -1 /gripper_ctrl std_msgs/Int16  \"data: 1000\"")
        rospy.sleep(1)

    def gripper_close(self):
        rospy.loginfo("Start close gripper")
        os.system("rosservice call /io_service/set_gripper_position \"{val: 10}\"")  # lebai原夹爪
        # os.system("rostopic pub -1 /gripper_ctrl std_msgs/Int16  \"data: 5\"")
        rospy.sleep(1)

    def shutdown(self):
        """关闭节点并清理资源"""
        moveit_commander.roscpp_shutdown()
        moveit_commander.os._exit(0)
        rospy.logwarn("Now shutting down grasp node...")

if __name__ == "__main__":
    MoveItGrasp()